use axum::{
    extract::{Path, State, Query, Json},
    http::{HeaderMap, StatusCode},
    response::{Response, IntoResponse},
    routing::{get, post, put, delete, any},
    Router,
};
use rustcloud_core::*;
use crate::client::{ServiceClient, ClientConfig};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// 网关配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GatewayConfig {
    pub port: u16,
    pub host: String,
    pub client_config: ClientConfig,
    pub routes: Vec<RouteConfig>,
}

impl Default for GatewayConfig {
    fn default() -> Self {
        Self {
            port: 8080,
            host: "0.0.0.0".to_string(),
            client_config: ClientConfig::default(),
            routes: Vec::new(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouteConfig {
    pub path: String,
    pub service_name: String,
    pub strip_prefix: bool,
    pub rewrite_path: Option<String>,
    pub allowed_methods: Vec<String>,
}

impl RouteConfig {
    pub fn new(path: impl Into<String>, service_name: impl Into<String>) -> Self {
        Self {
            path: path.into(),
            service_name: service_name.into(),
            strip_prefix: true,
            rewrite_path: None,
            allowed_methods: vec!["GET", "POST", "PUT", "DELETE"],
        }
    }
}

/// 网关状态
#[derive(Clone)]
pub struct GatewayState {
    pub client: Arc<ServiceClient>,
    pub routes: Arc<RwLock<Vec<RouteConfig>>>,
}

impl GatewayState {
    pub fn new(client: ServiceClient) -> Self {
        Self {
            client: Arc::new(client),
            routes: Arc::new(RwLock::new(Vec::new())),
        }
    }

    pub async fn add_route(&self, route: RouteConfig) {
        let mut routes = self.routes.write().await;
        // 检查是否已存在相同路径的路由
        routes.retain(|r| r.path != route.path);
        routes.push(route);
    }

    pub async fn remove_route(&self, path: &str) {
        let mut routes = self.routes.write().await;
        routes.retain(|r| r.path != path);
    }

    pub async fn get_routes(&self) -> Vec<RouteConfig> {
        let routes = self.routes.read().await;
        routes.clone()
    }

    pub async fn find_route(&self, path: &str) -> Option<RouteConfig> {
        let routes = self.routes.read().await;
        routes.iter().find(|r| path.starts_with(&r.path)).cloned()
    }
}

/// 网关服务器
pub struct GatewayServer {
    config: GatewayConfig,
    state: GatewayState,
}

impl GatewayServer {
    pub fn new(config: GatewayConfig) -> Self {
        let client = ServiceClient::new(config.client_config.clone());
        let state = GatewayState::new(client);
        Self { config, state }
    }

    pub async fn add_route(&self, route: RouteConfig) {
        self.state.add_route(route).await;
    }

    pub async fn remove_route(&self, path: &str) {
        self.state.remove_route(path).await;
    }

    pub async fn get_routes(&self) -> Vec<RouteConfig> {
        self.state.get_routes().await
    }

    pub fn router(&self) -> Router {
        Router::new()
            .route("/gateway/routes", get(list_routes))
            .route("/gateway/routes", post(add_route_handler))
            .route("/gateway/routes/*path", delete(remove_route_handler))
            .route("/gateway/health", get(health_check))
            .route("/*path", any(proxy_handler))
            .with_state(self.state.clone())
    }

    pub async fn serve(self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let addr = format!("{}", self.config.host, self.config.port);
        let router = self.router();

        println!("Gateway server starting on {}", addr);
        let listener = tokio::net::TcpListener::bind(addr).await?;
        axum::serve(listener, router).await?;
        Ok(())
    }
}

/// 路由列表处理器
async fn list_routes(
    State(state): State<GatewayState>
) -> Result<Json<Vec<RouteConfig>>, StatusCode> {
    let routes = state.get_routes().await;
    Ok(Json(routes))
}

/// 添加路由处理器
async fn add_route_handler(
    State(state): State<GatewayState>,
    Json(route): Json<RouteConfig>
) -> Result<Json<serde_json::Value>, StatusCode> {
    state.add_route(route).await;
    Ok(Json(serde_json::json!({"status": "success"})))
}

/// 删除路由处理器
async fn remove_route_handler(
    State(state): State<GatewayState>,
    Path(path): Path<String>
) -> Result<Json<serde_json::Value>, StatusCode> {
    state.remove_route(&path).await;
    Ok(Json(serde_json::json!({"status": "success"})))
}

/// 健康检查处理器
async fn health_check() -> Result<Json<serde_json::Value>, StatusCode> {
    Ok(Json(serde_json::json!({
        "status": "ok",
        "timestamp": chrono::Utc::now().to_rfc3339()
    })))
}

/// 代理处理器
async fn proxy_handler(
    State(state): State<GatewayState>,
    Path(path): Path<String>,
    headers: HeaderMap,
    Query(params): Query<HashMap<String, String>>,
    axum::body::Bytes(body): axum::body::Bytes
) -> Result<impl IntoResponse, StatusCode> {
    // 查找匹配的路由
    let route = state.find_route(&path).await.ok_or(StatusCode::NOT_FOUND)?;
    
    // 检查HTTP方法是否允许
    let method = headers.get("method").and_then(|v| v.to_str().ok()).unwrap_or("GET").to_uppercase();
    if !route.allowed_methods.iter().any(|m| m.to_uppercase() == method) {
        return Err(StatusCode::METHOD_NOT_ALLOWED);
    }

    // 构建目标路径
    let mut target_path = path.clone();
    if route.strip_prefix {
        target_path = path.replacen(&route.path, "", 1);
    }
    target_path = target_path.trim_start_matches('/').to_string();

    if let Some(rewrite) = &route.rewrite_path {
        target_path = rewrite.clone();
    }

    // 构建请求
    let mut request = Request::new(
        route.service_name.clone(),
        &method,
        &target_path
    );

    request.query_params = params;

    if !body.is_empty() {
        request.set_body(body.to_vec());
    }

    for (key, value) in headers.iter() {
        if let Ok(key_str) = key.to_string().parse::<String>() {
            if let Ok(value_str) = value.to_str() {
                request.add_header(key_str, value_str);
            }
        }
    }

    // 调用服务
    match state.client.call_service(request).await {
        Ok(response) => {
            let mut http_response = Response::builder()
                .status(response.status);

            for (key, value) in response.headers {
                http_response = http_response.header(&key, &value);
            }

            let body = response.body.unwrap_or_default();
            Ok(http_response.body(body).unwrap())
        }
        Err(e) => {
            tracing::error!("Service call failed: {}", e);
            match e {
                ServiceError::ServiceNotFound(_) => Err(StatusCode::NOT_FOUND),
                ServiceError::NoAvailableInstances(_) => Err(StatusCode::SERVICE_UNAVAILABLE),
                ServiceError::CircuitBreakerOpen(_) => Err(StatusCode::SERVICE_UNAVAILABLE),
                ServiceError::TimeoutError(_) => Err(StatusCode::REQUEST_TIMEOUT),
                _ => Err(StatusCode::INTERNAL_SERVER_ERROR),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::StatusCode;
    use tower::ServiceExt;

    #[tokio::test]
    async fn test_gateway_server() {
        let config = GatewayConfig::default();
        let mut server = GatewayServer::new(config);
        
        server.add_route(RouteConfig::new("/api/users", "user-service")).await;
        let routes = server.get_routes().await;
        assert_eq!(routes.len(), 1);
        assert_eq!(routes[0].path, "/api/users");
        assert_eq!(routes[0].service_name, "user-service");
    }

    #[tokio::test]
    async fn test_find_route() {
        let state = GatewayState::new(ServiceClient::new(ClientConfig::default()));
        state.add_route(RouteConfig::new("/api", "api-service")).await;
        state.add_route(RouteConfig::new("/api/users", "user-service")).await;

        let route = state.find_route("/api/users/123").await.unwrap();
        assert_eq!(route.service_name, "user-service");
        assert_eq!(route.path, "/api/users");
    }

    #[tokio::test]
    async fn test_health_check_endpoint() {
        let state = GatewayState::new(ServiceClient::new(ClientConfig::default()));
        let app = Router::new().route("/health", get(health_check)).with_state(state);

        let response = app
            .oneshot(axum::http::Request::builder()
                .uri("/health")
                .body(Body::empty())
                .unwrap())
            .await
            .unwrap();

        assert_eq!(response.status(), StatusCode::OK);
    }
}